import numpy as np
import cv2
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import pickle

def apply_filters(image):
    # 均值滤波
    mean_filtered = cv2.blur(image, (5,5))
    
    # 高斯滤波
    gaussian_filtered = cv2.GaussianBlur(image, (5,5), 0)
    
    # Sobel滤波
    sobel_x = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=3)
    sobel_y = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=3)
    sobel = np.sqrt(sobel_x**2 + sobel_y**2)
    
    # Canny边缘检测
    canny = cv2.Canny(image, 100, 200)
    
    return mean_filtered, gaussian_filtered, sobel, canny

def create_features(image):
    # 应用滤波器
    mean_filtered, gaussian_filtered, sobel, canny = apply_filters(image)
    
    # 将原图和所有滤波结果展平并组合
    features = np.column_stack((
        image.ravel(),
        mean_filtered.ravel(),
        gaussian_filtered.ravel(),
        sobel.ravel(),
        canny.ravel()
    ))
    
    return features

def main():
    # 读取图像
    image = cv2.imread('Sandstone_imgs/Sandstone_1.tif', cv2.IMREAD_GRAYSCALE)
    segment = cv2.imread('Sandstone_imgs/Sandstone_1_segment.tif', cv2.IMREAD_GRAYSCALE)
    
    print(f"图像形状: {image.shape}")
    
    # 创建特征
    X = create_features(image)
    y = segment.ravel()
    
    print("完成从砂岩截面图1及其对应分区中获取X和y")
    
    # 分割训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    print("完成train_test_split")
    
    # 训练随机森林模型
    clf = RandomForestClassifier(n_estimators=100, random_state=42)
    clf.fit(X_train, y_train)
    print("完成随机森林模型clf的训练")
    
    # 计算准确率
    accuracy = clf.score(X_test, y_test)
    print(f"准确率: {accuracy}")
    
    # 保存模型
    with open('sandstone_clf.pkl', 'wb') as f:
        pickle.dump(clf, f)
    print("已保存保存随机森林模型clf到硬盘")

if __name__ == "__main__":
    main() 